package com.zhima.util;

import cn.hutool.json.JSONUtil;
import com.aliyun.bailian20230601.Client;
import com.aliyun.bailian20230601.models.GetPromptRequest;
import com.aliyun.bailian20230601.models.GetPromptResponse;
import com.aliyun.broadscope.bailian.sdk.AccessTokenClient;
import com.aliyun.broadscope.bailian.sdk.ApplicationClient;
import com.aliyun.broadscope.bailian.sdk.models.CompletionsRequest;
import com.aliyun.broadscope.bailian.sdk.models.CompletionsResponse;
import com.aliyun.broadscope.bailian.sdk.models.ConnectOptions;
import com.aliyun.teaopenapi.models.Config;
import com.zhima.exception.ZhimaOjException;
import lombok.Data;

import java.util.HashMap;
import java.util.Map;

import static com.zhima.properties.AliyunBaiLianProperties.*;

public class AiUtil {
    public static String completionsWithParams(String wrongCode, String codeLanguage) {
        AccessTokenClient accessTokenClient = new AccessTokenClient(ACCESS_KEY_ID, ACCESS_KEY_SECRET, AGENT_KEY);
        String token = accessTokenClient.getToken();

        ApplicationClient client = ApplicationClient.builder()
                .token(token)
                .connectOptions(new ConnectOptions(30000, 60000, 60000))
                .build();
        CompletionsRequest request = new CompletionsRequest()
                .setAppId(APP_ID)
                //设置模型参数topP的值
                .setTopP(0.3)
                //设置prompt和历史上下文
                .setPrompt(getPrompt(wrongCode, codeLanguage))
                .setMessages(null)
                .setParameters(new CompletionsRequest.Parameter()
                        //设置模型参数topK
                        .setTopK(50)
                        //设置模型参数seed
                        .setSeed(2222)
                        //设置模型参数temperature
                        .setTemperature(0.7)
                        //设置最大内容token数
                        .setMaxTokens(2000)
                        //设置停止词
                        //设置内容返回结构为message
                        .setResultFormat("message")
                );
        CompletionsResponse response = client.completions(request);
        if (!response.isSuccess()) {
            throw new ZhimaOjException(String.format("failed to create completion, requestId: %s, code: %s, message: %s",
                    response.getRequestId(), response.getCode(), response.getMessage()));
        }
        return response.getData().getChoices().get(0).getMessage().getContent();
    }

    public static String getPrompt(String wrongCode, String codeLanguage) {
        try {
            Config config = new Config().setAccessKeyId(ACCESS_KEY_ID).setAccessKeySecret(ACCESS_KEY_SECRET).setEndpoint(ENDPOINT);
            Client client = new Client(config);

            Map<String, String> map = new HashMap<>();
            map.put("AgentKey", AGENT_KEY);
            map.put("PromptId", PROMPT_ID);
            map.put("Vars", JSONUtil.toJsonStr(new WrongCode(wrongCode, codeLanguage)));
            GetPromptResponse response = client.getPrompt(GetPromptRequest.build(map));
            // 模板内容
            System.out.println(response.getBody().getData().getPromptContent());
            return response.getBody().getData().getPromptContent();
        } catch (Exception e) {
            throw new ZhimaOjException("获取模板失败");
        }
    }
}

@Data
class WrongCode {
    private String wrongCode;
    private String codeLanguage;

    public WrongCode(String wrongCode, String codeLanguage) {
        this.wrongCode = wrongCode;
        this.codeLanguage = codeLanguage;
    }
}